package com.example.imserver.websocket;

import com.example.imserver.service.WebSocketService;
import com.example.imserver.util.JwtUtil;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.net.URI;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

@Component
public class MessageWebSocketHandler extends TextWebSocketHandler {
    
    @Autowired
    private WebSocketService webSocketService;
    
    // 维护session到userId的反向映射
    private final Map<WebSocketSession, Long> sessionToUserId = new ConcurrentHashMap<>();
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 从连接参数中获取userId
        URI uri = session.getUri();
        String query = uri != null ? uri.getQuery() : null;
        Long userId = null;
        
        if (query != null) {
            String[] params = query.split("&");
            for (String param : params) {
                String[] keyValue = param.split("=");
                if (keyValue.length == 2 && "userId".equals(keyValue[0])) {
                    try {
                        userId = Long.parseLong(keyValue[1]);
                        break;
                    } catch (NumberFormatException e) {
                        // 忽略
                    }
                }
            }
        }
        
        // 如果query中没有，尝试从token获取
        if (userId == null) {
            String token = session.getHandshakeHeaders().getFirst("token");
            if (token != null) {
                userId = JwtUtil.getUserIdFromToken(token);
            }
        }
        
        if (userId != null) {
            webSocketService.addSession(userId, session);
            sessionToUserId.put(session, userId);
            System.out.println("WebSocket连接建立: userId=" + userId);
        } else {
            session.close();
        }
    }
    
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 可以在这里处理客户端发送的消息
        System.out.println("收到WebSocket消息: " + message.getPayload());
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        Long userId = sessionToUserId.remove(session);
        if (userId != null) {
            webSocketService.removeSession(userId);
            System.out.println("WebSocket连接关闭: userId=" + userId);
        }
    }
}

